from scipy import sparse
import scipy
import learning
import pickle
import os

class FileWordsMatrix(object):
    def __init__(self):
        self.matrixFreq = sparse.lil_matrix((1, 1), dtype='int32')        
        self.categories = {}
        self.words = {}
        self.retraceWords = {}
        
        
    def loadMatrix(self, cats, words, files):
        oneLevelUp = os.listdir('..')        
        if 'dtreematrixFreq.csv' in oneLevelUp \
        and 'dtreeindexedCategories.ser' in oneLevelUp and 'dtreeindexedWords.ser'  in oneLevelUp\
        and 'dtreeretraceWords.ser' in oneLevelUp:
            print 'Load from file !!'
            self.loadFromFiles()
        else:
            print 'Build Cats'
            for c in cats.bag.items():
                cId = c[1][0]
                cFreq = c[1][1]
                cName = c[0]
                self.categories[cId] = [cName, cFreq]
            print 'Build Words'
            for w in words.bag.items():
                wId = w[1][0]
                wFreq = w[1][1]
                wName = w[0]
                self.words[wId] = [wName, wFreq]
    
            self.matrixFreq = sparse.lil_matrix((files.bag.__len__(), words.bag.__len__()+1), dtype='int32')           
            print 'Build Matrix'
            lastCol = words.bag.__len__()
            for f in files.bag.items():
                cId = cats.bag[f[1][2]][0]
                fId = f[1][0]
                self.matrixFreq[fId, lastCol] = cId   
                for w in f[1][1].items():
                    wName = w[0]
                    if words.bag.has_key(wName) :
                        wId = words.bag[wName][0]
                        wFreq = w[1]
                        self.matrixFreq[fId, wId] = wFreq       
            
    def pruneMatrix(self, min_var):
        oneLevelUp = os.listdir('..')        
        if not('dtreematrixFreq.csv' in oneLevelUp \
        and 'dtreeindexedCategories.ser' in oneLevelUp and 'dtreeindexedWords.ser'  in oneLevelUp\
        and 'dtreeretraceWords.ser' in oneLevelUp):
            vars = scipy.zeros(self.matrixFreq.shape[1]-1)
            for i in range(self.matrixFreq.shape[1]-1):
                vars[i] = self.matrixFreq[:, i].toarray().var()

            notPrune = []
            for i in range(self.matrixFreq.shape[1]-1):
                if vars[i] > min_var:
                    self.retraceWords[i] = notPrune.__len__()                
                    notPrune.append(i)
           
            self.retraceWords[self.matrixFreq.shape[1]-1] = notPrune.__len__()                
            notPrune.append(self.matrixFreq.shape[1]-1)         
            self.matrixFreq = self.matrixFreq[:, notPrune]     
        print 'End of pruning'
    def saveDic(self, filename, dic):
        f = open('../' + filename, 'w')
        pickle.dump(dic, f)
        f.close()
        del f
    
    
    def loadDic(self,filename):
        f = open('../'+filename, 'r')
        dic = pickle.load(f)
        f.close()        
        del f        
        return dic


    def save(self):
        learning.saveMatrix(self.matrixFreq, '../'+'dtreematrixFreq.csv')
        self.saveDic('dtreeindexedCategories.ser', self.categories)
        self.saveDic('dtreeindexedWords.ser', self.words)
        self.saveDic('dtreeretraceWords.ser', self.retraceWords)
        
        
    def loadFromFiles(self):
        self.categories = self.loadDic('dtreeindexedCategories.ser')
        self.words = self.loadDic('dtreeindexedWords.ser')
        self.retraceWords = self.loadDic('dtreeretraceWords.ser')               
        self.matrixFreq = learning.openMatrixFile('../', 'dtreematrixFreq.csv')
